import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = nn.NLLLoss()

    def forward(self, input_x, labels):
        m = torch.nn.LogSoftmax(dim=1)
        out = self.op(m(input_x), labels)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    x = torch.FloatTensor(np.ones([3,5]).astype(np.float32))
    # pytorch only support LongTensor -> Int64 not match to Mindspore support int32
    labels = torch.LongTensor(np.random.random_sample([3]).astype(np.int32))
    # labels = torch.tensor([1, 0, 4])

    result = net(x, labels)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)


